import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.ticker import FuncFormatter
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Backtest:
    def __init__(self, initial_capital=1000000, commission=0.0003, slippage=0.0001):
        """
        初始化回测系统
        :param initial_capital: 初始资金
        :param commission: 交易佣金率
        :param slippage: 滑点率
        """
        self.initial_capital = initial_capital
        self.commission = commission
        self.slippage = slippage
    
    def run(self, df, strategy):
        """
        运行回测
        :param df: 包含交易信号的DataFrame
        :param strategy: 交易策略
        :return: 回测结果DataFrame
        """
        if df.empty:
            logger.warning("数据为空，无法运行回测")
            return pd.DataFrame()
        
        try:
            # 复制数据
            results = df.copy()
            
            # 初始化资金和持仓
            results['capital'] = self.initial_capital
            results['position'] = 0
            results['holdings'] = 0
            results['cash'] = self.initial_capital
            results['total'] = self.initial_capital
            
            # 执行回测
            for i in range(1, len(results)):
                # 获取上一交易日的持仓和现金
                prev_position = results.iloc[i-1]['position']
                prev_cash = results.iloc[i-1]['cash']
                
                # 获取当前交易日的信号和价格
                signal = results.iloc[i]['signal']
                close_price = results.iloc[i]['close']
                
                # 计算交易数量
                trade_shares = 0
                
                if signal == 1 and prev_position <= 0:  # 买入信号
                    # 计算可买入的最大股数（考虑佣金和滑点）
                    available_cash = prev_cash * 0.95  # 使用95%的可用资金
                    price_with_cost = close_price * (1 + self.slippage)
                    max_shares = int(available_cash / (price_with_cost * (1 + self.commission)))
                    
                    if max_shares > 0:
                        trade_shares = max_shares - prev_position
                        cost = trade_shares * price_with_cost * (1 + self.commission)
                        
                        # 更新资金和持仓
                        results.loc[results.index[i], 'position'] = max_shares
                        results.loc[results.index[i], 'cash'] = prev_cash - cost
                        results.loc[results.index[i], 'holdings'] = max_shares * close_price
                
                elif signal == -1 and prev_position >= 0:  # 卖出信号
                    if prev_position > 0:
                        trade_shares = -prev_position  # 全部卖出
                        price_received = close_price * (1 - self.slippage)
                        proceeds = trade_shares * price_received * (1 - self.commission)
                        
                        # 更新资金和持仓
                        results.loc[results.index[i], 'position'] = 0
                        results.loc[results.index[i], 'cash'] = prev_cash - proceeds
                        results.loc[results.index[i], 'holdings'] = 0
                else:  # 保持持仓不变
                    results.loc[results.index[i], 'position'] = prev_position
                    results.loc[results.index[i], 'cash'] = prev_cash
                    results.loc[results.index[i], 'holdings'] = prev_position * close_price
                
                # 计算总资产
                results.loc[results.index[i], 'total'] = results.iloc[i]['cash'] + results.iloc[i]['holdings']
            
            # 计算每日收益率
            results['daily_return'] = results['total'].pct_change()
            
            # 计算累积收益率
            results['cumulative_return'] = (1 + results['daily_return']).cumprod() - 1
            
            logger.info(f"{strategy.name}回测完成，最终资产: {results.iloc[-1]['total']:.2f}")
            return results
        except Exception as e:
            logger.error(f"回测运行失败: {e}")
            return pd.DataFrame()
    
    def evaluate(self, results, strategy, benchmark=None):
        """
        评估回测结果
        :param results: 回测结果DataFrame
        :param strategy: 交易策略
        :param benchmark: 基准数据DataFrame
        :return: 评估指标字典
        """
        if results.empty:
            logger.warning("回测结果为空，无法进行评估")
            return {}
        
        try:
            # 计算评估指标
            metrics = {}
            
            # 总交易日数
            total_days = len(results)
            
            # 累积收益率
            cumulative_return = results.iloc[-1]['cumulative_return']
            annual_return = (1 + cumulative_return) ** (252 / total_days) - 1
            
            # 最大回撤
            results['peak'] = results['total'].cummax()
            results['drawdown'] = (results['total'] - results['peak']) / results['peak']
            max_drawdown = results['drawdown'].min()
            
            # 夏普比率
            risk_free_rate = 0.03  # 假设无风险收益率为3%
            daily_risk_free = (1 + risk_free_rate) ** (1/252) - 1
            excess_returns = results['daily_return'] - daily_risk_free
            sharpe_ratio = np.sqrt(252) * excess_returns.mean() / excess_returns.std()
            
            # 胜率
            trades = results[results['signal'] != 0]
            winning_trades = trades[trades['daily_return'] > 0]
            win_rate = len(winning_trades) / len(trades) if len(trades) > 0 else 0
            
            # 盈亏比
            if len(winning_trades) > 0 and len(trades[trades['daily_return'] < 0]) > 0:
                avg_profit = winning_trades['daily_return'].mean()
                avg_loss = trades[trades['daily_return'] < 0]['daily_return'].mean()
                profit_loss_ratio = -avg_profit / avg_loss
            else:
                profit_loss_ratio = 0
            
            # 交易频率
            trade_frequency = len(trades) / total_days
            
            # 填充评估指标
            metrics['strategy'] = strategy.name
            metrics['total_days'] = total_days
            metrics['cumulative_return'] = cumulative_return
            metrics['annual_return'] = annual_return
            metrics['max_drawdown'] = max_drawdown
            metrics['sharpe_ratio'] = sharpe_ratio
            metrics['win_rate'] = win_rate
            metrics['profit_loss_ratio'] = profit_loss_ratio
            metrics['trade_frequency'] = trade_frequency
            metrics['final_capital'] = results.iloc[-1]['total']
            
            # 计算相对基准的表现
            if benchmark is not None and not benchmark.empty:
                # 确保基准数据与回测数据日期对齐
                benchmark = benchmark.reindex(results.index)
                
                # 计算基准收益率
                benchmark['daily_return'] = benchmark['close'].pct_change()
                benchmark['cumulative_return'] = (1 + benchmark['daily_return']).cumprod() - 1
                
                # 计算相对收益率
                relative_return = cumulative_return - benchmark.iloc[-1]['cumulative_return']
                metrics['relative_return'] = relative_return
                
                # 计算信息比率
                tracking_error = (results['daily_return'] - benchmark['daily_return']).std()
                if tracking_error > 0:
                    information_ratio = relative_return / (tracking_error * np.sqrt(252))
                    metrics['information_ratio'] = information_ratio
            
            logger.info(f"{strategy.name}评估完成: 年化收益率={annual_return:.2%}, 最大回撤={max_drawdown:.2%}, 夏普比率={sharpe_ratio:.2f}")
            return metrics
        except Exception as e:
            logger.error(f"回测评估失败: {e}")
            return {}
    
    def plot_results(self, results, strategy, benchmark=None, filename=None):
        """
        可视化回测结果
        :param results: 回测结果DataFrame
        :param strategy: 交易策略
        :param benchmark: 基准数据DataFrame
        :param filename: 保存图像的文件名
        """
        if results.empty:
            logger.warning("回测结果为空，无法生成图表")
            return
        
        try:
            plt.figure(figsize=(14, 10))
            
            # 绘制累积收益率
            plt.subplot(2, 1, 1)
            plt.plot(results.index, results['cumulative_return'], label=f'{strategy.name}收益率', color='blue')
            
            if benchmark is not None and not benchmark.empty:
                # 确保基准数据与回测数据日期对齐
                benchmark = benchmark.reindex(results.index)
                
                # 计算基准累积收益率
                benchmark['daily_return'] = benchmark['close'].pct_change()
                benchmark['cumulative_return'] = (1 + benchmark['daily_return']).cumprod() - 1
                
                plt.plot(benchmark.index, benchmark['cumulative_return'], label='沪深300收益率', color='red')
            
            plt.title(f'{strategy.name}回测结果 - 累积收益率')
            plt.xlabel('日期')
            plt.ylabel('收益率')
            plt.grid(True)
            plt.legend()
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
            plt.gcf().autofmt_xdate()
            
            # 绘制回撤
            plt.subplot(2, 1, 2)
            plt.fill_between(results.index, results['drawdown'], 0, color='red', alpha=0.3)
            plt.title(f'{strategy.name}回测结果 - 回撤')
            plt.xlabel('日期')
            plt.ylabel('回撤比例')
            plt.grid(True)
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
            plt.gcf().autofmt_xdate()
            
            plt.tight_layout()
            
            if filename:
                plt.savefig(filename)
                logger.info(f"回测结果图表已保存至: {filename}")
            
            plt.show()
        except Exception as e:
            logger.error(f"回测结果可视化失败: {e}")    